#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-09-11


import threading
from queue import Queue
import numpy as np
import time
import os
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
import argparse
import platform
from spirecv.algorithm.utils import calc_fov, calc_los_pos
import cv2
from mmpretrain.apis import ImageClassificationInferencer


class MMPreTrainInferNode_Cuda(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.launch_next_emit = self.get_param("launch_next_emit", True)
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.specified_output_topic = self.get_param("specified_output_topic", "")
        self.realtime_det = self.get_param("realtime_det", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.dataset_name = self.get_param("dataset_name", "imagenet")
        self.model_config = self.get_param("model_config", "/home/amov/spirecv-training-pro/mmpretrain/configs/vision_transformer/vit-base-p16_64xb64_in1k-384px.py")
        self.model_weight = self.get_param("model_weight", "/home/amov/Downloads/vit-base-p16_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-98e8652b.pth")
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        if self.model_config.startswith("sms::"):
            self.local_model_config = download_model(self.__class__.__name__, self.model_config)
            assert self.local_model_config is not None
        else:
            self.local_model_config = self.model_config

        if self.model_weight.startswith("sms::"):
            self.local_model_weight = download_model(self.__class__.__name__, self.model_weight)
            assert self.local_model_weight is not None
        else:
            self.local_model_weight = self.model_weight

        input_url = '/' + job_name + '/sensor/image_raw'
        if len(self.specified_input_topic) > 0:
            input_url = self.specified_input_topic

        output_url = '/' + job_name + '/classifier/results'
        if len(self.specified_output_topic) > 0:
            output_url = self.specified_output_topic

        calib_url = '/' + job_name + '/sensor/calibration_info'

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self.calib_width, self.calib_height = -1, -1
        self.camera_matrix = [712.12, 0,645.23, 0, 705.87, 327.34, 0, 0, 1]
        self.camera_matrix = np.array(self.camera_matrix).reshape(3, 3)
        self.distortion = [0.0, 0.0, 0.0, 0.0, 0.0]
        self.distortion = np.array(self.distortion)

        self._image_reader = Subscriber(
            input_url, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )
        self._calibration_reader = Subscriber(
            calib_url, 'sensor_msgs::CameraCalibration', self.calibration_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            output_url, 'std_msgs::Null',
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        self.inferencer = ImageClassificationInferencer(self.model_config, pretrained=self.model_weight, classes=[str(i) for i in range(1000)])
        self.correct_top1 = 0
        self.correct_top5 = 0

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._calibration_reader.kill()
        self._result_writer.kill()
        self._next_writer.kill()

    def image_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        img = sms2cvimg(msg)
        self.job_queue.put({'msg': msg, 'img': img})

    def calibration_callback(self, msg):
        self.calib_width = msg['width']
        self.calib_height = msg['height']

        self.camera_matrix = np.array(msg['K']).reshape(3, 3)
        self.distortion = np.array(msg['D'])

    def run(self):
        while self.is_running():
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break

            msg, img = msg_dict['msg'], msg_dict['img']

            t1 = time.time()
            file_name = msg['file_name'] if 'file_name' in msg else ''

            results = self.inferencer(img, show=False)[0]

            top1_cls = results['pred_label']
            top5_cls = np.argsort(results['pred_scores'])[-5:].tolist()
            res_msg = def_msg('std_msgs::Null')
            res_msg['top1'] = top1_cls
            res_msg['top5'] = top5_cls
            # probs = results[0].probs.data.cpu().numpy().tolist()
            self._result_writer.publish(res_msg)

            if top1_cls == msg['class_id']:
                self.correct_top1 += 1
            
            if msg['class_id'] in top5_cls:
                self.correct_top5 += 1

            if msg['img_id'] % 100 == 0:
                print("img_id = {}".format(msg['img_id']))

            if 'img_total' in msg and self.launch_next_emit:
                next_msg = def_msg('std_msgs::Boolean')
                next_msg['data'] = True
                self._next_writer.publish(next_msg)
            
                if 'img_total' in msg and msg['img_id'] + 1 == msg['img_total']:
                    print("Top1 Acc: {}".format(self.correct_top1 / msg['img_total']))
                    print("Top5 Acc: {}".format(self.correct_top5 / msg['img_total']))
                    self.emit_shutdown()

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = MMPreTrainInferNode_Cuda(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
